package com.example.portlocalforward.sshmanage;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import javax.annotation.PreDestroy;

import org.springframework.stereotype.Component;

import com.example.portlocalforward.controller.model.PortLocalForwardResponse;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@Component
public class SessionAndPortLocalForwardContainer {
    private static final ConcurrentHashMap<SessionKey, Session> SESSIONS = new ConcurrentHashMap<>();

    public void setStatus(List<PortLocalForwardResponse> portLocalForwardResponses) {
        closeUnUsedSession();
        // 动态设置状态
        Set<SessionAndPortLocalForward> sessionAndPortLocalForwardSet = new HashSet<>(
                getSessionAndPortLocalForwardList());

        for (PortLocalForwardResponse response : portLocalForwardResponses) {
            int status = sessionAndPortLocalForwardSet.contains(
                    new SessionAndPortLocalForward(
                            response.getSshTunnelIp(),
                            response.getSshTunnelPort(),
                            response.getSshTunnelUsername(),
                            response.getLocalPort(),
                            response.getRemoteHost(),
                            response.getRemotePort())) ? 1 : 0;
            response.setStatus(status);
        }
    }

    public List<String[]> add(List<PortLocalForwardResponse> portLocalForwardResponses) {
        removeUnConnectedSession();
        List<SessionAndPortLocalForward> isRunningSessionAndPortLocalForwardList = getSessionAndPortLocalForwardList();
        // 当前正在运行的端口本地转发列表
        Set<SessionAndPortLocalForward> isRunningSessionAndPortLocalForwardSet = new HashSet<>(
                isRunningSessionAndPortLocalForwardList);
        // 当前正在运行的本地端口列表
        Set<Integer> isRunningLocalPortSet = isRunningSessionAndPortLocalForwardList.stream()
                .map(SessionAndPortLocalForward::getLocalPort).collect(Collectors.toSet());
        // 创建失败的SessionKey列表
        Set<SessionKey> createErroredSessionKeys = new HashSet<>();
        // 创建失败的结果信息
        List<String[]> createErrorResult = new ArrayList<>();
        for (PortLocalForwardResponse portLocalForwardResponse : portLocalForwardResponses) {
            // 已经运行的就跳过
            if (isRunningSessionAndPortLocalForwardSet
                    .contains(SessionAndPortLocalForward.of(portLocalForwardResponse))) {
                continue;
            }
            // 本地端口已经存在就跳过，但是要记录错误信息
            if (isRunningLocalPortSet.contains(portLocalForwardResponse.getLocalPort())) {
                createErrorResult.add(new String[] { portLocalForwardResponse.getName(),
                        String.format("本地端口 %d 已启动", portLocalForwardResponse.getLocalPort()) });
                continue;
            }
            SessionKey sessionKey = SessionKey.of(portLocalForwardResponse);
            Session session = SESSIONS.get(sessionKey);
            if (session == null) {
                try {
                    if (createErroredSessionKeys.contains(sessionKey)) {
                        createErrorResult
                                .add(new String[] { portLocalForwardResponse.getName(), String.format("隧道连接失败") });
                        continue;
                    }
                    session = createSession(sessionKey);
                    SESSIONS.put(sessionKey, session);
                } catch (JSchException e) {
                    log.error("创建Session失败: {}", sessionKey, e);
                    createErrorResult.add(new String[] { portLocalForwardResponse.getName(), String.format("隧道连接失败") });
                    createErroredSessionKeys.add(sessionKey);
                    continue;
                }
            }
            try {
                session.setPortForwardingL(portLocalForwardResponse.getLocalPort(),
                        portLocalForwardResponse.getRemoteHost(), portLocalForwardResponse.getRemotePort());
                isRunningLocalPortSet.add(portLocalForwardResponse.getLocalPort());
            } catch (JSchException e) {
                log.error("添加端口本地转发失败: {}", portLocalForwardResponse, e);
                createErrorResult.add(new String[] { portLocalForwardResponse.getName(), String.format("添加端口本地转发失败") });
            }
        }
        return createErrorResult;
    }

    public void remove(List<PortLocalForwardResponse> portLocalForwardResponses) {
        removeUnConnectedSession();
        // 当前正在运行的端口本地转发列表
        Set<SessionAndPortLocalForward> sessionAndPortLocalForwardSet = new HashSet<>(
                getSessionAndPortLocalForwardList());
        // 需要停止的端口本地转发列表
        List<PortLocalForwardResponse> needStopPortLocalForwardResponses = portLocalForwardResponses.stream()
                .filter(response -> sessionAndPortLocalForwardSet.contains(SessionAndPortLocalForward.of(response)))
                .collect(Collectors.toList());
        for (PortLocalForwardResponse response : needStopPortLocalForwardResponses) {
            try {
                SessionKey sessionKey = SessionKey.of(response);
                Session session = SESSIONS.get(sessionKey);
                if (session != null && session.isConnected()) {
                    session.delPortForwardingL(response.getLocalPort());
                }
            } catch (JSchException e) {
                log.error("停止端口本地转发失败: {}", response, e);
            }
        }
        closeUnUsedSession();
    }

    // 判断是否有运行中的端口本地转发
    public boolean isRunning(List<PortLocalForwardResponse> portLocalForwardResponses) {
        Set<SessionAndPortLocalForward> sessionAndPortLocalForwardSet = new HashSet<>(
                getSessionAndPortLocalForwardList());
        for (PortLocalForwardResponse response : portLocalForwardResponses) {
            if (sessionAndPortLocalForwardSet.contains(SessionAndPortLocalForward.of(response))) {
                return true;
            }
        }
        return false;
    }

    private void removeUnConnectedSession() {
        Iterator<SessionKey> iterator = SESSIONS.keySet().iterator();
        while (iterator.hasNext()) {
            SessionKey sessionKey = iterator.next();
            Session session = SESSIONS.get(sessionKey);
            if (session == null || !session.isConnected()) {
                iterator.remove();
                continue;
            }
        }
    }

    private void closeUnUsedSession() {
        Iterator<SessionKey> iterator = SESSIONS.keySet().iterator();
        while (iterator.hasNext()) {
            SessionKey sessionKey = iterator.next();
            Session session = SESSIONS.get(sessionKey);
            if (session == null || !session.isConnected()) {
                iterator.remove();
                continue;
            }
            try {
                String[] forwardings = session.getPortForwardingL();
                if (forwardings == null || forwardings.length == 0) {
                    session.disconnect();
                    iterator.remove();
                }
            } catch (JSchException e) {
                session.disconnect();
                iterator.remove();
            }
        }
    }

    private Session createSession(SessionKey sessionKey) throws JSchException {
        JSch jsch = new JSch();
        Session session = jsch.getSession(sessionKey.getUsername(),
                sessionKey.getHost(), sessionKey.getPort());
        session.setPassword(sessionKey.getPassword());

        // 设置不验证主机密钥（生产环境建议验证）
        java.util.Properties config = new java.util.Properties();
        config.put("StrictHostKeyChecking", "no");
        session.setConfig(config);
        // 连接
        session.connect(30000); // 5秒超时
        return session;
    }

    private List<SessionAndPortLocalForward> getSessionAndPortLocalForwardList() {
        List<SessionAndPortLocalForward> result = new ArrayList<>();

        // 遍历所有Session
        SESSIONS.forEach((sessionKey, session) -> {
            if (session == null || !session.isConnected()) {
                return;
            }
            try {
                String[] forwardings = session.getPortForwardingL();
                if (forwardings == null || forwardings.length == 0) {
                    return;
                }

                // 遍历每个转发配置，解析本地端口转发信息
                for (String forwarding : forwardings) {
                    // 格式："lport:host:rport"，按":"分割
                    String[] parts = forwarding.split(":");
                    if (parts.length >= 3) {
                        int localPort = Integer.parseInt(parts[0]);
                        String remoteHost = parts[1];
                        int remotePort = Integer.parseInt(parts[2]);

                        // 创建SessionAndPortLocalForward对象
                        SessionAndPortLocalForward item = new SessionAndPortLocalForward();
                        item.setHost(sessionKey.getHost());
                        item.setPort(sessionKey.getPort());
                        item.setUsername(sessionKey.getUsername());
                        item.setLocalPort(localPort);
                        item.setRemoteHost(remoteHost);
                        item.setRemotePort(remotePort);
                        result.add(item);
                    }
                }
            } catch (JSchException e) {
                // 获取端口转发信息失败，跳过该Session
            }
        });

        return result;
    }

    /**
     * 应用关闭时清理所有 SSH sessions 和端口转发
     * 防止端口占用问题
     */
    @PreDestroy
    public void destroy() {
        log.info("应用关闭，开始清理所有 SSH sessions 和端口转发...");
        Iterator<SessionKey> iterator = SESSIONS.keySet().iterator();
        while (iterator.hasNext()) {
            SessionKey sessionKey = iterator.next();
            Session session = SESSIONS.get(sessionKey);
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
            iterator.remove();
        }
        log.info("所有 SSH sessions 和端口转发已清理完成");
    }
}
